import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import datetime
from Mr_Zhong.utils.log import Logger
from sklearn.metrics import roc_auc_score, roc_curve, classification_report
from sklearn.metrics import mean_squared_error, mean_absolute_error
from sklearn.metrics import confusion_matrix
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import seaborn as sns
import joblib


class PowerLoadModel(object):
    def __init__(self, filename):
        # 配置日志记录
        logfile_name = "predict_" + datetime.datetime.now().strftime('%Y%m%d%H%M%S')
        self.logfile = Logger('../', logfile_name).get_logger()
        # 获取数据源
        self.data_source = pd.read_csv(filename, encoding='utf-8')

def feature_engineering(data, logger):
    logger.info('===============开始进行特征工程处理===============')
    result = data.copy(deep=True)
    logger.info("===============开始处理测试集数据特征===================")
    # 1.提取出特征

    # X_train = result.loc[:, ['Age', 'BusinessTravel', 'Department', 'DistanceFromHome', 'Education',
    #                          'EducationField', 'EnvironmentSatisfaction', 'Gender', 'JobInvolvement', 'JobLevel',
    #                          'JobRole', 'JobSatisfaction', 'MaritalStatus', 'MonthlyIncome', 'NumCompaniesWorked',
    #                          'OverTime', 'PercentSalaryHike', 'PerformanceRating', 'RelationshipSatisfaction',
    #                          'StockOptionLevel', 'TotalWorkingYears', 'TrainingTimesLastYear',
    #                          'WorkLifeBalance', 'YearsAtCompany', 'YearsInCurrentRole', 'YearsSinceLastPromotion',
    #                          'YearsWithCurrManager']]

    # XGBoost的数据
    X_train = result.loc[:, ['Age', 'BusinessTravel', 'Department', 'DistanceFromHome', 'Education',
                             'EducationField', 'EnvironmentSatisfaction', 'Gender', 'JobInvolvement', 'JobLevel',
                             'JobSatisfaction', 'MaritalStatus', 'MonthlyIncome', 'NumCompaniesWorked',
                             'OverTime', 'PercentSalaryHike', 'RelationshipSatisfaction', 'StockOptionLevel', 'TotalWorkingYears', 'TrainingTimesLastYear',
                              'YearsAtCompany', 'YearsInCurrentRole', 'YearsSinceLastPromotion', 'YearsWithCurrManager']]
    Y_train = result.iloc[:, -1]

    # 1.1修改BusinessTravel中的值, Non-Travel,Travel_Rarely,Travel_Frequently
    travel_BusinessTravel_map = {
        'Non-Travel': 0,
        'Travel_Rarely': 1,
        'Travel_Frequently': 2
    }
    X_train['BusinessTravel'] = X_train['BusinessTravel'].map(travel_BusinessTravel_map)

    # 1.2修改Department中的值, Human Resources, Research & Development , Sales
    travel_Department_map = {
        'Human Resources': 1,
        'Research & Development': 2,
        'Sales': 3
    }
    X_train['Department'] = X_train['Department'].map(travel_Department_map)

    # 1.3修改EducationField中的值, Life Sciences , Medical , Marketing, Technical Degree, Other, Human Resources
    travel_EducationField_map = {
        'Life Sciences': 1,
        'Medical': 2,
        'Marketing': 3,
        'Technical Degree': 4,
        'Other': 5,
        'Human Resources': 6
    }
    X_train['EducationField'] = X_train['EducationField'].map(travel_EducationField_map)

    # 1.4修改Gender中的值, Male,female
    travel_Gender_map = {
        'Male': 1,
        'Female': 0,
    }
    X_train['Gender'] = X_train['Gender'].map(travel_Gender_map)

    # # 1.5修改JobRole中的值, Sales Executive ,Research Scientist,Laboratory Technician,Manufacturing Director ,
    # # Healthcare Representative ,Manager ,Sales Representative  ,Research Director ,Human Resources
    # travel_JobRole_map = {
    #     'Sales Executive': 1,
    #     'Research Scientist': 2,
    #     'Laboratory Technician': 3,
    #     'Manufacturing Director': 4,
    #     'Healthcare Representative': 5,
    #     'Manager': 6,
    #     'Sales Representative': 7,
    #     'Research Director': 8,
    #     'Human Resources': 9
    # }
    # X_train['JobRole'] = X_train['JobRole'].map(travel_JobRole_map)

    # 1.6修改MaritalStatus中的值, Married ,Single ,Divorced
    travel_MaritalStatus_map = {
        'Divorced': 0,
        'Single': 1,
        'Married': 2
    }
    X_train['MaritalStatus'] = X_train['MaritalStatus'].map(travel_MaritalStatus_map)

    # 1.7修改OverTime中的值, No, Yes
    travel_OverTime_map = {
        'No': 0,
        'Yes': 1
    }
    X_train['OverTime'] = X_train['OverTime'].map(travel_OverTime_map)
    logger.info("===============测试集数据特征化结束===================")

    # X_train.to_csv('X_test.csv', index=False)


    return X_train, Y_train

def plot_1(y_test, y_pre):
    # 计算ROC曲线的假正例率（FPR）、真正例率（TPR）和阈值
    fpr, tpr, thresholds = roc_curve(y_test, y_pre)
    # 计算AUC分数
    auc_score = roc_auc_score(y_test, y_pre)

    # 创建一个新的图形，设置图形大小
    plt.figure(figsize=(8, 6))
    # 绘制ROC曲线，并标注AUC值
    plt.plot(fpr, tpr, label=f'ROC Curve (AUC = {auc_score:.2f})')
    # 绘制随机猜测的ROC曲线，作为参考
    plt.plot([0, 1], [0, 1], 'k--', label='Random Guessing')
    # 设置x轴标签
    plt.xlabel('False Positive Rate (FPR)')
    # 设置y轴标签
    plt.ylabel('True Positive Rate (TPR)')
    # 设置图形标题
    plt.title('Receiver Operating Characteristic (ROC) Curve')
    # 显示图例
    plt.legend(loc='lower right')
    # 启用网格
    plt.grid(True)
    # 保存图片
    plt.savefig('../result/预测效果.png')
    # 显示图形
    # plt.show()


def plot_confusion_matrix(y_true, y_pre):
    """
    绘制混淆矩阵
    :param y_true: 真实标签
    :param y_pred: 预测标签
    :param save_path: 保存路径
    """
    # 计算混淆矩阵
    cm = confusion_matrix(y_true, y_pre)
    # 使用 seaborn 可视化
    plt.figure(figsize=(6, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', cbar=False)
    plt.title('Confusion Matrix')
    plt.xlabel('Predicted Label')
    plt.ylabel('True Label')

    # 保存图片
    plt.savefig('../result/混淆矩阵.png')
    # 显示图像
    # plt.show()


def plot_classification_report(y_true, y_pred, cmap='Blues'):
    """
    可视化分类报告
    :param y_true: 真实标签
    :param y_pred: 预测标签
    :param cmap:   颜色映射
    """
    # 获取分类报告字典
    report = classification_report(y_true, y_pred, output_dict=True)

    # 转换为 DataFrame，并转置以便显示
    df_report = pd.DataFrame(report).T  # shape: (n_classes + 3, 3) 包含 precision, recall, f1-score, support

    # 提取 support 并从主表中删除
    support = df_report['support']  # 所有类别的 support 数量
    df_report = df_report.drop('support', axis=1)  # 只保留 precision, recall, f1-score

    # 绘制 heatmap
    plt.figure(figsize=(8, 6))
    sns.heatmap(df_report, annot=True, fmt=".2f", cmap=cmap, cbar=False)

    # 在图下方标注 support
    for i, label in enumerate(df_report.index):
        plt.text(i + 0.5, len(df_report) + 0.3, f'supp={int(support[label])}',
                 ha='center', va='center', fontsize=10)

    plt.title("Classification Report Heatmap")
    plt.yticks(rotation=0)
    plt.tight_layout()
    plt.savefig('../result/评估报告.png')
    # plt.show()

if __name__ == '__main__':
    # 1.加载数据集
    input_file = os.path.join('../data', 'test2.csv')
    log1 = PowerLoadModel(input_file)
    # 2.特征工程
    x_test, y_test = feature_engineering(log1.data_source, log1.logfile)
    # 3.加载模型
    model = joblib.load('../model/xgb.pkl')
    log1.logfile.info("================== 开始加载模型并且预测 ==================")
    # 4.模型预测
    y_pre = model.predict_proba(x_test)[:, 1]
    y_pre1 = model.predict(x_test)
    # 5.预测结果评价
    print(f"模型在验证集上的auc:{roc_auc_score(y_test, y_pre)}")
    mse_test = mean_squared_error(y_true=y_test, y_pred=y_pre)
    mae_test = mean_absolute_error(y_true=y_test, y_pred=y_pre)
    print(f"模型在验证集上的均方误差：{mse_test}")
    print(f"模型在验证集上的平均绝对误差：{mae_test}")
    print(f"分类评估报告:{classification_report(y_test, y_pre1)}")
    log1.logfile.info(f"模型对新数据进行预测的均方误差：{mse_test}")
    log1.logfile.info(f"模型对新数据进行预测的平均绝对误差：{mae_test}")
    log1.logfile.info(f"模型对新数据进行预测的平均绝对误差：{roc_auc_score(y_test, y_pre)}")
    log1.logfile.info(f"模型对新数据进行预测的平均绝对误差：{roc_auc_score(y_test, y_pre)}")
    log1.logfile.info(f"模型对新数据进行预测的准确率(accuracy):{accuracy_score(y_test, y_pre1)}")
    log1.logfile.info(f"模型对新数据进行预测的精确率(precision):{precision_score(y_test, y_pre1, pos_label=0)}")
    log1.logfile.info(f"模型对新数据进行预测的召回率(recall):{recall_score(y_test, y_pre1, pos_label=0)}")
    log1.logfile.info(f"模型对新数据进行预测的f1分数:{f1_score(y_test, y_pre1, pos_label=0)}")
    print(f"分类评估报告:{classification_report(y_test, y_pre1)}")
    log1.logfile.info("================== 模型预测结束 ==================")
    # 6.可视化操作
    plot_1(y_test, y_pre)
    plot_classification_report(y_test, y_pre1)

